#ifndef __PE_READER__
#define __PE_READER__

#include "pefmt.h"
#include "archive.h"

namespace xm
{
	DEFINE_PROGID(WinPeParser);

	typedef enum
	{
		PPL_IMAGE_FILE				= 0x00000001,		
		PPL_IMAGE_HEADERS			= 0x00000002,		
		PPL_IMAGE_SECTIONS			= 0x00000004,

		PPL_IMAGE_BASIC				= PPL_IMAGE_FILE|PPL_IMAGE_HEADERS|PPL_IMAGE_SECTIONS,

		PPL_IMAGE_IMPORT			= 0x00000010,		
		PPL_IMAGE_EXPORT			= 0x00000020,		
		PPL_IMAGE_RESOURCE			= 0x00000040,	
		PPL_IMAGE_TLS				= 0x00000080,			
		PPL_IMAGE_DEBUG				= 0x00000100,		
		PPL_IMAGE_RELOCATION		= 0x00000200,	
		PPL_IMAGE_DELAYIMPORT		= 0x00000400,	
		PPL_IMAGE_LOADCONFIG		= 0x00000800,	
		PPL_IMAGE_CERTIFICATE		= 0x00001000,	
		PPL_IMAGE_BOUNDIMPORT		= 0x00002000,
		PPL_IMAGE_CLRRUNTIME		= 0x00004000,

		PPL_IMAGE_ALL				= 0xFFFFFFFF,
	}WINPE_PARSE_LEVEL;

	typedef enum
	{
		PET_IMAGE_NUL				= 0x00000000,
		PET_IMAGE_X86				= 0x00000001,		
		PET_IMAGE_X64				= 0x00000002,		
		PET_IMAGE_CUI				= 0x00000004,		
		PET_IMAGE_GUI				= 0x00000008,		
		PET_IMAGE_SYS				= 0x00000010,		
		PET_IMAGE_EXE				= 0x00000020,		
		PET_IMAGE_DLL				= 0x00000040,		
	}WINPE_IMAGE_TYPE;

	template<typename TX32, typename TX64>
	union TXBound
	{
		TX32	X32;
		TX64	X64;
	};

	class WinPeTypes : public IKArchive
	{
	public:
		//////////////////////////////////////////////////////////////////////////
		typedef IMAGE_DOS_HEADER						TImgDosHdr;	
		typedef IMAGE_FILE_HEADER						TImgFileHdr;

		typedef IMAGE_NT_HEADERS32						TImgNtHdr32;	
		typedef IMAGE_NT_HEADERS64						TImgNtHdr64;
		typedef TXBound<TImgNtHdr32, TImgNtHdr64>		TImgNtHdrBnd;

		typedef IMAGE_OPTIONAL_HEADER32					TImgOptHdr32;
		typedef IMAGE_OPTIONAL_HEADER64					TImgOptHdr64;
		typedef TXBound<TImgOptHdr32, TImgOptHdr64>		TImgOptHdrBnd;	


		typedef	IMAGE_SECTION_HEADER					TImgSecHdr;
		typedef IMAGE_DATA_DIRECTORY					TImgDatDir;

		typedef IMAGE_THUNK_DATA32						TImgThkDat32;
		typedef IMAGE_THUNK_DATA64						TImgThkDat64;
		typedef TXBound<TImgThkDat32, TImgThkDat64> 	TImgThkDatBnd;

		typedef IMAGE_IMPORT_DESCRIPTOR					TImgImpDes;
		typedef IMAGE_EXPORT_DIRECTORY					TImgExpDir;
		typedef IMAGE_RESOURCE_DIRECTORY				TImgResDir;

		typedef IMAGE_TLS_DIRECTORY32					TImgTlsDir32;	
		typedef IMAGE_TLS_DIRECTORY64					TImgTlsDir64;
		typedef TXBound<TImgTlsDir32, TImgTlsDir64> 	TImgTlsDirBnd;

		typedef IMAGE_DEBUG_DIRECTORY					TImgDbgDir;
		typedef IMAGE_LOAD_CONFIG_DIRECTORY32			TImgCfgDir32;
		typedef IMAGE_LOAD_CONFIG_DIRECTORY64			TImgCfgDir64;
		typedef TXBound<TImgCfgDir32, TImgCfgDir64>		TImgCfgDirBnd;

		typedef IMAGE_DELAY_IMPORT_DESCRIPTOR			TImgDlyDes;
		typedef WIN_CERTIFICATE							TImgCrtDir;
		typedef IMAGE_BOUND_IMPORT_DESCRIPTOR			TImgBndDes;
		typedef IMAGE_BOUND_FORWARDER_REF				TImgBndFwd;
		typedef IMAGE_COR20_HEADER						TImgClrDir;

		typedef IMAGE_RESOURCE_DIRECTORY_ENTRY			TImgResEnt;
		typedef IMAGE_RESOURCE_DATA_ENTRY				TImgResDat;
		typedef IMAGE_BASE_RELOCATION					TImgBasRel;
	};

	//[uuid("019B1050-DE1F-4387-A8B3-78844E833BAB")]

	class IPeReader : public WinPeTypes
	{
	public:
		 
	public:
		typedef UINT64			QWORD;
		typedef IKFileStream	TStream;

		virtual TStream* Stream() = 0;
	
		//////////////////////////////////////////////////////////////////////////
		virtual BOOL IsValidRaw(DWORD dwRaw, DWORD dwSize = 1) = 0;
		virtual BOOL IsValidRva(DWORD dwRva, DWORD dwSize = 1) = 0;
		virtual BOOL IsValidVa(QWORD dwVa, DWORD dwSize = 1) = 0;

		virtual BOOL Va2Rva(QWORD dwVa, DWORD dwSize = 1, DWORD* lpRva = NULL) = 0;
		virtual BOOL Rva2Va(DWORD dwRva, DWORD dwSize = 1, QWORD* lpVa = NULL) = 0;

		virtual BOOL Rva2Raw(DWORD dwRva, DWORD dwSize = 1, DWORD* lpRaw = NULL, DWORD* lpSec = NULL) = 0;
		virtual BOOL Raw2Rva(DWORD dwRaw, DWORD dwSize = 1, DWORD* lpRva = NULL, DWORD* lpSec = NULL) = 0;

		virtual BOOL ReadDatRaw(DWORD dwRaw, LPVOID lpBuff, DWORD dwLen) = 0;
		virtual BOOL ReadDatRva(DWORD dwRva, LPVOID lpBuff, DWORD dwLen) = 0;
		virtual BOOL ReadDatVa(QWORD dwVa, LPVOID lpBuff, DWORD dwLen) = 0;

		virtual DWORD ReadDatRaw2(DWORD dwRaw, LPVOID lpBuff, DWORD dwLen) = 0;
		virtual DWORD ReadDatRva2(DWORD dwRva, LPVOID lpBuff, DWORD dwLen) = 0;

	
		//////////////////////////////////////////////////////////////////////////
		virtual BOOL FileSize(QWORD* lpSize) = 0;	

		//////////////////////////////////////////////////////////////////////////
		virtual BOOL IsWow64() = 0;
		virtual BOOL ImageType(DWORD dwMask) = 0;
		virtual BOOL ImageBaseVa(QWORD* lpImgVa = NULL, DWORD* lpSize = NULL) = 0;	
		virtual BOOL ImageChecksum(DWORD* lpChksum = NULL) = 0;
		virtual BOOL ImageTimeStamp(DWORD* lpTimeStamp) = 0;
		virtual BOOL ImageHeaderSize(DWORD* lpSize = NULL) = 0;
		virtual BOOL EntryPointRva(DWORD* lpRva = NULL) = 0;			

		//////////////////////////////////////////////////////////////////////////
		virtual TImgDosHdr* DosHeader(DWORD* lpSize = NULL)  = 0;
		virtual TImgNtHdrBnd* NtHeaders(DWORD* lpSize = NULL) = 0;
		virtual TImgOptHdrBnd* OptionalHeader(DWORD* lpSize = NULL) = 0;
		virtual TImgFileHdr* FileHeader(DWORD* lpSize = NULL) = 0;		
		virtual BOOL DataDirectoryCount(DWORD* lpCount = NULL) = 0;
		virtual TImgDatDir* DataDirectory(DWORD dwIndex = 0, DWORD* lpSize = NULL) = 0;

		//////////////////////////////////////////////////////////////////////////
		virtual BOOL AlignSectionAboveRva(DWORD dwRva, DWORD* lpOut) = 0;
		virtual BOOL AlignSectionBelowRva(DWORD dwRva, DWORD* lpOut) = 0;
		virtual BOOL AlignFileAboveRaw(DWORD dwRaw, DWORD* lpOut) = 0;	
		virtual BOOL AlignFileBelowRaw(DWORD dwRaw, DWORD* lpOut) = 0;
		virtual BOOL AlignRelocAboveRva(DWORD dwRva, DWORD* lpOut) = 0;	
		virtual BOOL AlignRelocBelowRva(DWORD dwRva, DWORD* lpOut) = 0;	

		//////////////////////////////////////////////////////////////////////////
		virtual TImgSecHdr* SectionHeader(DWORD dwIndex = 0, DWORD* lpSize = NULL) = 0;
		virtual BOOL SectionCount(DWORD* lpCount = NULL) = 0;		
		virtual BOOL SectionRva(DWORD dwIndex, DWORD* lpRva = NULL, DWORD* lpSize = NULL) = 0;	
		virtual BOOL SectionRaw(DWORD dwIndex, DWORD* lpRaw = NULL, DWORD* lpSize = NULL) = 0;
		virtual LPCSTR SectionName(DWORD dwIndex) = 0;	

		//////////////////////////////////////////////////////////////////////////
		virtual BOOL OverlayRaw(DWORD* lpRaw = NULL, DWORD* lpSize = NULL)	= 0;

	
		//////////////////////////////////////////////////////////////////////////
		// 01. Export Data Directory
		virtual TImgExpDir* ExportDirectory(TU32* lpSize = NULL) = 0;
		virtual LPCSTR ExportModule() = 0;	
		virtual BOOL IsExportModule(LPCSTR szModule) = 0;	
		virtual LPCSTR ExportFunctionNext(LPCSTR szLastFunc = NULL) = 0;
		virtual BOOL IsExportFunction(LPCSTR szFunction) = 0;
		virtual BOOL ExportFunctionRva(LPCSTR szFuncName, DWORD* lpRva = NULL) = 0;	

		//////////////////////////////////////////////////////////////////////////
		// 02. Import Data Directory
		virtual TImgImpDes* ImportDirectory(LPCSTR szModule, TU32* lpSize = NULL) = 0;
		virtual BOOL ImportCount(DWORD* lpLibCount = NULL, DWORD* lpFuncCount = NULL) = 0;	
		virtual LPCSTR NextImportModule(LPCSTR szLastModu = NULL) = 0;	
		virtual BOOL IsImportFunction(LPCSTR szModule, LPCSTR szFunction) = 0;
		virtual LPCSTR NextImportFunction(LPCSTR szModule, LPCSTR szLastFunc = NULL) = 0;
		virtual BOOL ImportFunctionSlotRva(LPCSTR szModule, LPCSTR szFuncName, DWORD* lpRva = NULL) = 0;
		virtual BOOL ImportFunctionNameRVA(LPCSTR szModule, LPCSTR szFuncName, DWORD* lpRva = NULL) = 0;
		//////////////////////////////////////////////////////////////////////////
		// 03. Resource Data Directory
		virtual BOOL ResourceRva(LPCSTR szType, LPCSTR szNameId, LPCSTR szLang, DWORD* lpRva = NULL, DWORD* lpSize = NULL) = 0;
		virtual LPCSTR NextResourceType(LPCSTR szLast = NULL) = 0;
		virtual LPCSTR NextResourceNameId(LPCSTR szType, LPCSTR szLast = NULL)	= 0;
		virtual LPCSTR NextResourceLanguage(LPCSTR szType, LPCSTR szNameId, LPCSTR szLast = NULL) = 0;

		virtual BOOL ResourceCount(DWORD* lpCount = NULL) = 0;
		virtual BOOL ResourceTypeCount(DWORD* lpSize = NULL) = 0;
		virtual BOOL ResourceNameIdCount(LPCSTR szType, DWORD* lpSize) = 0;
		virtual BOOL ResourceLanguageCount(LPCSTR szType, LPCSTR szNameId, DWORD* lpCount = NULL) = 0;

		virtual TImgResDir* ResourceTypeDirectory(DWORD* lpSize = NULL) = 0;
		virtual TImgResDir* ResourceNameIdDirectory(LPCSTR szType, DWORD* lpSize = NULL) = 0;
		virtual TImgResDir* ResourceLanguageDirectory(LPCSTR szType, LPCSTR szNameId, DWORD* lpSize = NULL) = 0;

		//////////////////////////////////////////////////////////////////////////
		// 04. Exception Table

		//////////////////////////////////////////////////////////////////////////
		// 05. Certificate Table
		virtual TImgCrtDir* CertificateDirectory(TU32* lpSize = NULL) = 0;
		virtual BOOL CertificateData(DWORD* lpRaw = NULL, DWORD* lpSize = NULL) = 0;	

		//////////////////////////////////////////////////////////////////////////
		// 06. Base Relocation Table
		virtual BOOL IsRelocRva(DWORD dwRva) = 0;
		virtual BOOL RelocDelta(QWORD dwNewBase, LONG* lpDelta = NULL) = 0;
		virtual BOOL RelocBlock(LONG lDelta, DWORD dwBuffRva, LPVOID lpBuff, DWORD dwSize) = 0;

		//////////////////////////////////////////////////////////////////////////
		// 07. Debug Data Directory
		virtual TImgDbgDir* DebugDirectory(DWORD* lpSize = NULL) = 0;
		virtual LPCSTR DebugFilePath() = 0;

		//////////////////////////////////////////////////////////////////////////
		// 08. Architecture (Reserved)

		//////////////////////////////////////////////////////////////////////////
		// 09. Global Ptr (Reserved)

		//////////////////////////////////////////////////////////////////////////
		// 10. TLS Table
		virtual TImgTlsDirBnd* TlsDirectory(TU32* lpSize = NULL) = 0;
		virtual BOOL TlsTempleteRva(DWORD* lpRva = NULL, DWORD* lpSize = NULL) = 0;
		virtual BOOL TlsCallBackCount(DWORD* lpCount = NULL) = 0;
		virtual BOOL TlsCallBackRva(DWORD dwIndex, DWORD* lpRva = NULL) = 0;

		//////////////////////////////////////////////////////////////////////////
		// 11. Load Config Table
		virtual TImgCfgDirBnd* LoadConfigDirectory(TU32* lpSize = NULL) = 0;
		virtual BOOL SehHandlerCount(DWORD* lpCount = NULL) = 0;
		virtual BOOL SehHandlerRva(DWORD dwIndex, DWORD* lpRva = NULL) = 0;

		//////////////////////////////////////////////////////////////////////////
		// 12. Bound Import Table
		virtual TImgBndDes* BoundImportDescriptor(LPCSTR szModule, TU32* lpSize = NULL) = 0;
		virtual BOOL BoundModuleCount(DWORD* lpCount) = 0;
		virtual LPCSTR NextBoundModule(LPCSTR szLast = NULL) = 0;
		virtual BOOL BoundTimeStamp(LPCSTR szModule, DWORD* lpTms) = 0;
		virtual BOOL BoundForwardCount(LPCSTR szModule, DWORD* lpCount) = 0;
		virtual LPCSTR BoundForwardModule(LPCSTR szModule, DWORD nIndex) = 0;

		//////////////////////////////////////////////////////////////////////////
		// 13. Import Address Table
		virtual BOOL IsImportFunctionSlotRva(DWORD dwRva) = 0;
		virtual BOOL ImportAddresTable(DWORD* lpRva = NULL, DWORD* lpSize = NULL) = 0;
		
		//////////////////////////////////////////////////////////////////////////
		// 14. Delay Import Descriptor
		virtual TImgDlyDes* DelayImportDescriptor(LPCSTR szModule, TU32* lpSize = NULL) = 0;
		virtual BOOL DelayImportCount(TU32* lpCount = NULL) = 0;
		virtual LPCSTR NextDelayModule(LPCSTR szLast = NULL) = 0;
		virtual LPCSTR NextDelayFunction(LPCSTR szModule, LPCSTR szLast = NULL) = 0;
		virtual BOOL DelayFunctionSlotRva(LPCSTR szModule, LPCSTR szFuncName, DWORD* lpRva = NULL) = 0;

		//////////////////////////////////////////////////////////////////////////
		// 15. CLR Runtime Header 
		virtual TImgClrDir* ClrRuntimeDirectory(TU32* lpSize = NULL) = 0;
	};//IPeReader

	// {F2886A0C-4F17-4886-84E7-CA8C966FA267}
	DEFINE_IID(IPeReader, "{F2886A0C-4F17-4886-84E7-CA8C966FA267}",
		0xf2886a0c, 0x4f17, 0x4886, 0x84, 0xe7, 0xca, 0x8c, 0x96, 0x6f, 0xa2, 0x67);


};//WinPe

#endif